import torch
from torch import nn


class VectorConstruction(nn.Module):
    """
    向量构造层，将输入的隐藏层维度 d_model 转换为每个注意力头的维度 d_k * num_heads
    """

    def __init__(self, d_model: int, heads: int, d_k: int, bais: bool = True):
        super().__init__()
        self.heads = heads
        self.d_k = d_k
        self.linear = nn.Linear(d_model, d_k * heads, bias=bais)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [seq_len, batch_size, d_model] or [batch_size, d_model]
        # 保存输入的维度信息
        head_shape = x.shape[:-1]

        # 线性变换
        x = self.linear(x)
        # 维度变换
        # [seq_len, batch_size, d_k * heads] -> [seq_len, batch_size, heads, d_k]
        x = x.view(*head_shape, self.heads, self.d_k)
        return x


class MultiHeadAttention(nn.Module):
    def __init__(
        self, d_model: int, heads: int, dropout_prob: float = 0.1, bais: bool = True
    ):
        super().__init__()
        # 注意力头的数量
        self.heads = heads
        # 输入隐藏层维度
        self.d_model = d_model
        # 每个注意力头的维度，拼在一起就是隐藏层维度
        self.d_k = d_model // heads

        # 构建Q、K、V向量
        self.query = VectorConstruction(self.d_model, self.heads, self.d_k, bais)
        self.key = VectorConstruction(self.d_model, self.heads, self.d_k, bais)
        self.value = VectorConstruction(self.d_model, self.heads, self.d_k, bais)

        self.softmax = nn.Softmax(dim=-1)

        self.output = nn.Linear(d_model, d_model)

        self.dropout = nn.Dropout(dropout_prob)

        self.scale = 1 / math.sqrt(self.d_k)

        # 注意力权重，保存下来用于分析
        self.attn = None
    
    def get_score(self, query: torch.Tensor, key: torch.Tensor) -> torch.Tensor:
        """
        计算注意力分数
        """
        # query: [seq_len, batch_size, heads, d_k]
        # key: [seq_len, batch_size, heads, d_k]
        # scores: [seq_len, seq_len, batch_size, heads]
        # 计算点积
        scores = torch.einsum('ibhd,jbhd->ijbh', query, key)
        return scores
    
    def prepare_mask(self, mask: torch.Tensor, query_shape: List[int], key_shape: List[int]):
        """
        准备注意力掩码
        """
        # mask 维度（query_len, key_len, batch_size, 1）
        assert mask.shape[0] == 1 or mask.shape[0] == query_shape[0]
        assert mask.shape[1] == key_shape[0]
        assert mask.shape[2] == 1 or mask.shape[2] == query_shape[1]

        # 增加头维度
        mask = mask.unsqueeze(-1)
        return mask
    
    def forward(self, *, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, mask: Optional[torch.Tensor] = None):
        """
        前向传播
        """
        # query: [seq_len, batch_size, heads, d_k]
        # key: [seq_len, batch_size, heads, d_k]
        # scores: [seq_len, seq_len, batch_size, heads]
        # mask: [query_len, key_len, batch_size, 1]
        query = self.query(query)
        key = self.key(key)
        value = self.value(value)

        if mask is not None:
            mask = self.prepare_mask(mask, query.shape, key.shape)

        # 计算注意力分数
        scores = self.get_score(query, key)
        scores = scores * self.scale

        # 应用掩码
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))

        # 计算注意力权重
        attn = self.softmax(scores)
        # 应用dropout
        attn = self.dropout(attn)

        # 计算输出
        # attn: [batch_size, num_heads, seq_len, seq_len]
        # value: [batch_size, seq_len, num_heads, d_k]
        # x: [batch_size, seq_len, num_heads, d_k]
        x = torch.einsum('ijbh,jbhd->ibhd', attn, value)

        self.attn = attn.detach()

        # 维度变换
        # [batch_size, seq_len, num_heads, d_k] -> [batch_size, seq_len, d_model]
        x = x.reshape(x.shape[0], x.shape[1], -1)

        # 输出层
        output = self.output(x)
        return output


